{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uhYG8yG47rKW"
      },
      "source": [
        "# License\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\")\n",
        "```\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License.\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qWACTvqSLHNo"
      },
      "source": [
        "# Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "suU32_cvK_-h"
      },
      "outputs": [],
      "source": [
        "# Uncomment to install the covid_vhh_design package\n",
        "\n",
        "# !pip install git+https://github.com/google-research/google-research.git#subdirectory=covid_vhh_design"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a51RdViEciv2"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JAFVmV4j7-hX"
      },
      "outputs": [],
      "source": [
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import seaborn as sns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C2xfG8K58fou"
      },
      "outputs": [],
      "source": [
        "from covid_vhh_design import covid"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ct5kxN1dH0iV"
      },
      "outputs": [],
      "source": [
        "pd.set_option('display.width', 200)\n",
        "pd.set_option('display.max_colwidth', None)\n",
        "pd.set_option('display.max_rows', 200)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2FA3lDyFcrRr"
      },
      "source": [
        "# Loading data files\n",
        "Annotated AlphaSeq measurements of the three described libraries are stored in\n",
        "* `data/round1.csv`: The 1st library\n",
        "* `data/round2.csv`: The 2nd library\n",
        "* `data/round3.csv`: The 3rd library\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UVj40rpF8tZS"
      },
      "outputs": [],
      "source": [
        "df = covid.load_df('round1.csv.gz')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "92gSsCTpdVaZ"
      },
      "source": [
        "# Description of columns\n",
        "All three CSV files have the following columns:\n",
        "* `source_key`: A unique identifier of an assayed VHH sequence. Note that there can be sources with different keys but the same sequence if multiple sequences replicates were assayed.\n",
        "* `target_key`: A unique identifier of a CoV RBD binding target.\n",
        "* `replica`: The AlphaSeq experiment replica. round0 has 3 replicas (1-3), round1 and round2 six (1-6).\n",
        "* `value`: An AlphaSeq log KD measurement. There exist exactly one value per `source_key x target_key x replica` triplet. A value can be infinity (`inf`) if the binding strength is below the detecable threshold.\n",
        "\n",
        "The remaining `target_*` and `source_*` columns store annotations of targets and sources respectively, which will be described in the following."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vlr2jD6Jde6Y"
      },
      "outputs": [],
      "source": [
        "df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zVDVFrw-gpf3"
      },
      "source": [
        "## Target columns\n",
        "* `target_name`: A unique name per `target_key`\n",
        "  - `MATalpha_Control_[1-6]`: Positive control targets\n",
        "  - `MATalpha_Negative_[1-3]`: Negative control targets\n",
        "  - `SARS-CoV1_RBD`: Wildtype SARS-CoV1\n",
        "  - `SARS-CoV2_RBD`: Wildtype SARS-CoV2\n",
        "  - `SARS-CoV2_RBD_*`: SARS-CoV2 mutants\n",
        "  - All other: CoV-related targets\n",
        "* `target_is_control`: Whether a target is a control\n",
        "* `target_control_type`: Whether a target is 'positive' or 'negative' a control\n",
        "* `target_group`: The name of a target group. E.g. 'CoV2' for wildtype CoV2 or CoV2 mutants"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dvYKbUV39OnC"
      },
      "outputs": [],
      "source": [
        "df.loc[:, df.columns.str.startswith('target_')].drop_duplicates()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A9DexOq6iq5_"
      },
      "source": [
        "## Source columns\n",
        "* `source_name`: A unique name per `source_key`.\n",
        "* `source_is_control`: Whether a source is a control. Control sources were included for santiy-checking and should bind only weakly or not at all.\n",
        "* `source_control_type`: Whether a source is 'positive' or 'negative' a control.\n",
        "* `source_num_mutation`: The number of mutations (Hamming distance) from the wildtype sequence.\n",
        "* `source_category`: One of the following categories\n",
        "    - `mutant`: A mutant sequence of VHH-72 (`source_num_mutations \u003e 0`).\n",
        "    - `parent`: The parent sequence VHH-72 (`source_num_mutations == 0`).\n",
        "    - `negative_control`: A negative control.\n",
        "    - `positive_control`: A positive control.\n",
        "    - `special`: A 'special' sequences that was not derived from VHH-72, e.g. 'SARS_VHH44' or 'MERS_VHH55'.\n",
        "* `source_hash`: The MD5 hash of `source_seq`.\n",
        "* `source_replica`: The sequence replica. Distiguishes sources with the exact same sequence. For example, round0 contains 46 copies of the parent sequence VHH-72 (1-46).\n",
        "* `source_distance`: Same as `source_num_mutations`, except that distances are all `\u003e0` instead of `\u003e=0`, or `NaN`.\n",
        "* `source_score`: The negative log-likelihood of the VAE that was used for scoring mutants.\n",
        "* `source_group`: Describes how sequences were designed per round.\n",
        "* `source_seq`: The amino-acid sequence.\n",
        "* `source_std_group`: Standardized annotation for sequence design across rounds.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V3RYbIaPvk9j"
      },
      "outputs": [],
      "source": [
        "source_df = df.loc[:, df.columns.str.startswith('source_')].drop_duplicates()\n",
        "source_df.head()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6dbWYLgrFT3"
      },
      "source": [
        "Round0 sequences were grouped as follows:\n",
        "* `mbo`: Model-designed sequences,\n",
        "  - `cdrh12_multies1`: Sequences with mutations in CDR1 or CDR2 designed by optimizing the 1st` VAE.\n",
        "  - `cdrh12_multies2: Sequences with mutations in CDR1 or CDR2 designed by optimizing the 2nd` VAE.\n",
        "  - `cdrh12_random`: Sequences with mutations in CDR1 or CDR2 scored randomly.\n",
        "  - `cdr3_multies1`: Sequences with mutations in CDR3 designed by optimizing the 1st VAE.\n",
        "  - `cdr3_multies2`: Sequences with mutations in CDR3 designed by optimizing the 2nd VAE.\n",
        "  - `cdr3_random`: Sequences with mutations in CDR3 obtained by reservoir sampling.\n",
        "  - `mutant`: Uncategorized sequences with randomly sampled mutations.\n",
        "* `negative_control`: Negative control sources.\n",
        "* `parent`: The parent sequence VHH-72.\n",
        "* `positive_control`: Positive control sources.\n",
        "* `singles`: A single-mutant sequence of VHH-72.\n",
        "* `special`: A 'special' sequences that was not derived from VHH-72, e.g. 'SARS_VHH44' or 'MERS_VHH55'."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fCg7e6V1rAN3"
      },
      "outputs": [],
      "source": [
        "source_df.groupby('source_std_group')['source_group'].unique()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9B4LViwinyVT"
      },
      "outputs": [],
      "source": [
        "# Round0 contains 46 copies of the parent sequence\n",
        "source_df.query('source_num_mutations == 0').head()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bX8_QQRLGuqB"
      },
      "outputs": [],
      "source": [
        "# The number of sources with N mutations aways from VHH-72\n",
        "source_df.value_counts('source_num_mutations')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vUDjrJnnnZ9l"
      },
      "outputs": [],
      "source": [
        "# Sequences replicas have different `source_replica` IDs.\n",
        "source_df.query('source_num_mutations == 0')['source_replica'].nunique()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AKm8BRWElXfD"
      },
      "outputs": [],
      "source": [
        "# Special sequences\n",
        "source_df.query('source_category == \"special\"')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4eWq4cErZC_R"
      },
      "source": [
        "# Visualization of AlphaSeq measuremnts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5rUbRzGVyovL"
      },
      "source": [
        "AlphaSeq log KD measurements are stored in the `value` column. A value can be infinity (`inf`) if the binding strength is below the detecable threshold. The following functions illustrate how to plot 1) the distribution of non-infinity values per target, and 2) the percentage of infinity values (non-binding events) per target."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jA5OuCUSH8F4"
      },
      "outputs": [],
      "source": [
        "def plot_non_inf_values_per_target(df):\n",
        "  data = df.loc[~np.isinf(df['value'])]\n",
        "  _, ax = plt.subplots(figsize=(15, 6))\n",
        "  order = data.groupby('target_name')['value'].mean().sort_values().index\n",
        "  sns.boxplot(\n",
        "      data=data,\n",
        "      x='target_name',\n",
        "      order=order,\n",
        "      y='value',\n",
        "      ax=ax)\n",
        "  ax.set_ylabel('AlphaSeq log KD')\n",
        "  ax.figure.canvas.draw()\n",
        "  ax.set_xticklabels(ax.get_xticklabels(), rotation=30, ha='right')\n",
        "\n",
        "\n",
        "plot_non_inf_values_per_target(df)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OSdJnxLkswnd"
      },
      "outputs": [],
      "source": [
        "def plot_percentage_of_inf_values_per_target(df):\n",
        "  data = df.copy()\n",
        "  data['value'] = np.isinf(data['value'])\n",
        "  data = (\n",
        "      data.groupby('target_name')['value'].mean().reset_index()\n",
        "      .sort_values('value'))\n",
        "\n",
        "  _, ax = plt.subplots(figsize=(15, 6))\n",
        "  sns.barplot(\n",
        "      data=data,\n",
        "      x='target_name',\n",
        "      y='value',\n",
        "      ax=ax)\n",
        "  ax.set_ylabel('% Inf values')\n",
        "  ax.figure.canvas.draw()\n",
        "  ax.set_xticklabels(ax.get_xticklabels(), rotation=30, ha='right')\n",
        "\n",
        "\n",
        "plot_percentage_of_inf_values_per_target(df)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
